{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Meridian MLflow Demo\n",
        "\n",
        "This Colab demonstrates how to integrate Meridian model runs with [MLflow tracking](https://mlflow.org/docs/latest/ml/tracking/). We'll cover the necessary installation and setup steps, followed by a simple example that showcases Meridian's MLflow configuration and data retrieval."
      ],
      "metadata": {
        "id": "rM3_e1OWy7RO"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GicRPam0mUhF"
      },
      "source": [
        "<a name=\"install\"></a>\n",
        "## Step 0: Install"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pDdX9WofM2fx"
      },
      "source": [
        "1\\. Make sure you are using one of the available GPU Colab runtimes which is **required** to run Meridian. You can change your notebook's runtime in `Runtime > Change runtime type` in the menu. All users can use the T4 GPU runtime which is sufficient to run the demo colab, free of charge. Users who have purchased one of Colab's paid plans have access to premium GPUs (such as V100, A100 or L4 Nvidia GPU)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nFYRTDuesa1P"
      },
      "source": [
        "2\\. Install the latest version of Meridian with MLflow dependencies (MLflow support is available as of version 1.1.3), and verify that GPU is available."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h1jAk386jF3k",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d7cb9a77-2050-4153-fde2-c7525cdf295a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: google-meridian[and-cuda,colab,mlflow] in /usr/local/lib/python3.11/dist-packages (1.1.3)\n",
            "Requirement already satisfied: arviz in /usr/local/lib/python3.11/dist-packages (from google-meridian[and-cuda,colab,mlflow]) (0.21.0)\n",
            "Requirement already satisfied: altair>=5 in /usr/local/lib/python3.11/dist-packages (from google-meridian[and-cuda,colab,mlflow]) (5.5.0)\n",
            "Requirement already satisfied: immutabledict in /usr/local/lib/python3.11/dist-packages (from google-meridian[and-cuda,colab,mlflow]) (4.2.1)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.11/dist-packages (from google-meridian[and-cuda,colab,mlflow]) (1.5.1)\n",
            "Requirement already satisfied: natsort<8,>=7.1.1 in /usr/local/lib/python3.11/dist-packages (from google-meridian[and-cuda,colab,mlflow]) (7.1.1)\n",
            "Requirement already satisfied: numpy<3,>=2.0.2 in /usr/local/lib/python3.11/dist-packages (from google-meridian[and-cuda,colab,mlflow]) (2.0.2)\n",
            "Requirement already satisfied: pandas<3,>=2.2.2 in /usr/local/lib/python3.11/dist-packages (from google-meridian[and-cuda,colab,mlflow]) (2.2.2)\n",
            "Requirement already satisfied: scipy<2,>=1.13.1 in /usr/local/lib/python3.11/dist-packages (from google-meridian[and-cuda,colab,mlflow]) (1.15.3)\n",
            "Requirement already satisfied: tensorflow<2.19,>=2.18 in /usr/local/lib/python3.11/dist-packages (from google-meridian[and-cuda,colab,mlflow]) (2.18.0)\n",
            "Requirement already satisfied: tensorflow-probability<0.26,>=0.25 in /usr/local/lib/python3.11/dist-packages (from google-meridian[and-cuda,colab,mlflow]) (0.25.0)\n",
            "Requirement already satisfied: tf-keras<2.19,>=2.18 in /usr/local/lib/python3.11/dist-packages (from google-meridian[and-cuda,colab,mlflow]) (2.18.0)\n",
            "Requirement already satisfied: xarray in /usr/local/lib/python3.11/dist-packages (from google-meridian[and-cuda,colab,mlflow]) (2025.3.1)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.11/dist-packages (from google-meridian[and-cuda,colab,mlflow]) (5.9.5)\n",
            "Requirement already satisfied: python-calamine in /usr/local/lib/python3.11/dist-packages (from google-meridian[and-cuda,colab,mlflow]) (0.3.2)\n",
            "Requirement already satisfied: mlflow in /usr/local/lib/python3.11/dist-packages (from google-meridian[and-cuda,colab,mlflow]) (3.1.0)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.11/dist-packages (from altair>=5->google-meridian[and-cuda,colab,mlflow]) (3.1.6)\n",
            "Requirement already satisfied: jsonschema>=3.0 in /usr/local/lib/python3.11/dist-packages (from altair>=5->google-meridian[and-cuda,colab,mlflow]) (4.24.0)\n",
            "Requirement already satisfied: narwhals>=1.14.2 in /usr/local/lib/python3.11/dist-packages (from altair>=5->google-meridian[and-cuda,colab,mlflow]) (1.44.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.11/dist-packages (from altair>=5->google-meridian[and-cuda,colab,mlflow]) (24.2)\n",
            "Requirement already satisfied: typing-extensions>=4.10.0 in /usr/local/lib/python3.11/dist-packages (from altair>=5->google-meridian[and-cuda,colab,mlflow]) (4.14.0)\n",
            "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.11/dist-packages (from pandas<3,>=2.2.2->google-meridian[and-cuda,colab,mlflow]) (2.9.0.post0)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.11/dist-packages (from pandas<3,>=2.2.2->google-meridian[and-cuda,colab,mlflow]) (2025.2)\n",
            "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.11/dist-packages (from pandas<3,>=2.2.2->google-meridian[and-cuda,colab,mlflow]) (2025.2)\n",
            "Requirement already satisfied: absl-py>=1.0.0 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (1.4.0)\n",
            "Requirement already satisfied: astunparse>=1.6.0 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (1.6.3)\n",
            "Requirement already satisfied: flatbuffers>=24.3.25 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (25.2.10)\n",
            "Requirement already satisfied: gast!=0.5.0,!=0.5.1,!=0.5.2,>=0.2.1 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (0.6.0)\n",
            "Requirement already satisfied: google-pasta>=0.1.1 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (0.2.0)\n",
            "Requirement already satisfied: libclang>=13.0.0 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (18.1.1)\n",
            "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (3.4.0)\n",
            "Requirement already satisfied: protobuf!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<6.0.0dev,>=3.20.3 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (5.29.5)\n",
            "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (2.32.3)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (75.2.0)\n",
            "Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (1.17.0)\n",
            "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (3.1.0)\n",
            "Requirement already satisfied: wrapt>=1.11.0 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (1.17.2)\n",
            "Requirement already satisfied: grpcio<2.0,>=1.24.3 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (1.73.0)\n",
            "Requirement already satisfied: tensorboard<2.19,>=2.18 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (2.18.0)\n",
            "Requirement already satisfied: keras>=3.5.0 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (3.8.0)\n",
            "Requirement already satisfied: h5py>=3.11.0 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (3.14.0)\n",
            "Requirement already satisfied: ml-dtypes<0.5.0,>=0.4.0 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (0.4.1)\n",
            "Requirement already satisfied: tensorflow-io-gcs-filesystem>=0.23.1 in /usr/local/lib/python3.11/dist-packages (from tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (0.37.1)\n",
            "Requirement already satisfied: decorator in /usr/local/lib/python3.11/dist-packages (from tensorflow-probability<0.26,>=0.25->google-meridian[and-cuda,colab,mlflow]) (4.4.2)\n",
            "Requirement already satisfied: cloudpickle>=1.3 in /usr/local/lib/python3.11/dist-packages (from tensorflow-probability<0.26,>=0.25->google-meridian[and-cuda,colab,mlflow]) (3.1.1)\n",
            "Requirement already satisfied: dm-tree in /usr/local/lib/python3.11/dist-packages (from tensorflow-probability<0.26,>=0.25->google-meridian[and-cuda,colab,mlflow]) (0.1.9)\n",
            "Requirement already satisfied: nvidia-cublas-cu12==12.5.3.2 in /usr/local/lib/python3.11/dist-packages (from tensorflow[and-cuda]<2.19,>=2.18; extra == \"and-cuda\"->google-meridian[and-cuda,colab,mlflow]) (12.5.3.2)\n",
            "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.5.82 in /usr/local/lib/python3.11/dist-packages (from tensorflow[and-cuda]<2.19,>=2.18; extra == \"and-cuda\"->google-meridian[and-cuda,colab,mlflow]) (12.5.82)\n",
            "Requirement already satisfied: nvidia-cuda-nvcc-cu12==12.5.82 in /usr/local/lib/python3.11/dist-packages (from tensorflow[and-cuda]<2.19,>=2.18; extra == \"and-cuda\"->google-meridian[and-cuda,colab,mlflow]) (12.5.82)\n",
            "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.5.82 in /usr/local/lib/python3.11/dist-packages (from tensorflow[and-cuda]<2.19,>=2.18; extra == \"and-cuda\"->google-meridian[and-cuda,colab,mlflow]) (12.5.82)\n",
            "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.5.82 in /usr/local/lib/python3.11/dist-packages (from tensorflow[and-cuda]<2.19,>=2.18; extra == \"and-cuda\"->google-meridian[and-cuda,colab,mlflow]) (12.5.82)\n",
            "Requirement already satisfied: nvidia-cudnn-cu12==9.3.0.75 in /usr/local/lib/python3.11/dist-packages (from tensorflow[and-cuda]<2.19,>=2.18; extra == \"and-cuda\"->google-meridian[and-cuda,colab,mlflow]) (9.3.0.75)\n",
            "Requirement already satisfied: nvidia-cufft-cu12==11.2.3.61 in /usr/local/lib/python3.11/dist-packages (from tensorflow[and-cuda]<2.19,>=2.18; extra == \"and-cuda\"->google-meridian[and-cuda,colab,mlflow]) (11.2.3.61)\n",
            "Requirement already satisfied: nvidia-curand-cu12==10.3.6.82 in /usr/local/lib/python3.11/dist-packages (from tensorflow[and-cuda]<2.19,>=2.18; extra == \"and-cuda\"->google-meridian[and-cuda,colab,mlflow]) (10.3.6.82)\n",
            "Requirement already satisfied: nvidia-cusolver-cu12==11.6.3.83 in /usr/local/lib/python3.11/dist-packages (from tensorflow[and-cuda]<2.19,>=2.18; extra == \"and-cuda\"->google-meridian[and-cuda,colab,mlflow]) (11.6.3.83)\n",
            "Requirement already satisfied: nvidia-cusparse-cu12==12.5.1.3 in /usr/local/lib/python3.11/dist-packages (from tensorflow[and-cuda]<2.19,>=2.18; extra == \"and-cuda\"->google-meridian[and-cuda,colab,mlflow]) (12.5.1.3)\n",
            "Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /usr/local/lib/python3.11/dist-packages (from tensorflow[and-cuda]<2.19,>=2.18; extra == \"and-cuda\"->google-meridian[and-cuda,colab,mlflow]) (2.21.5)\n",
            "Requirement already satisfied: nvidia-nvjitlink-cu12==12.5.82 in /usr/local/lib/python3.11/dist-packages (from tensorflow[and-cuda]<2.19,>=2.18; extra == \"and-cuda\"->google-meridian[and-cuda,colab,mlflow]) (12.5.82)\n",
            "Requirement already satisfied: matplotlib>=3.5 in /usr/local/lib/python3.11/dist-packages (from arviz->google-meridian[and-cuda,colab,mlflow]) (3.10.0)\n",
            "Requirement already satisfied: h5netcdf>=1.0.2 in /usr/local/lib/python3.11/dist-packages (from arviz->google-meridian[and-cuda,colab,mlflow]) (1.6.1)\n",
            "Requirement already satisfied: xarray-einstats>=0.3 in /usr/local/lib/python3.11/dist-packages (from arviz->google-meridian[and-cuda,colab,mlflow]) (0.9.1)\n",
            "Requirement already satisfied: mlflow-skinny==3.1.0 in /usr/local/lib/python3.11/dist-packages (from mlflow->google-meridian[and-cuda,colab,mlflow]) (3.1.0)\n",
            "Requirement already satisfied: Flask<4 in /usr/local/lib/python3.11/dist-packages (from mlflow->google-meridian[and-cuda,colab,mlflow]) (3.1.1)\n",
            "Requirement already satisfied: alembic!=1.10.0,<2 in /usr/local/lib/python3.11/dist-packages (from mlflow->google-meridian[and-cuda,colab,mlflow]) (1.16.2)\n",
            "Requirement already satisfied: docker<8,>=4.0.0 in /usr/local/lib/python3.11/dist-packages (from mlflow->google-meridian[and-cuda,colab,mlflow]) (7.1.0)\n",
            "Requirement already satisfied: graphene<4 in /usr/local/lib/python3.11/dist-packages (from mlflow->google-meridian[and-cuda,colab,mlflow]) (3.4.3)\n",
            "Requirement already satisfied: gunicorn<24 in /usr/local/lib/python3.11/dist-packages (from mlflow->google-meridian[and-cuda,colab,mlflow]) (23.0.0)\n",
            "Requirement already satisfied: pyarrow<21,>=4.0.0 in /usr/local/lib/python3.11/dist-packages (from mlflow->google-meridian[and-cuda,colab,mlflow]) (18.1.0)\n",
            "Requirement already satisfied: scikit-learn<2 in /usr/local/lib/python3.11/dist-packages (from mlflow->google-meridian[and-cuda,colab,mlflow]) (1.6.1)\n",
            "Requirement already satisfied: sqlalchemy<3,>=1.4.0 in /usr/local/lib/python3.11/dist-packages (from mlflow->google-meridian[and-cuda,colab,mlflow]) (2.0.41)\n",
            "Requirement already satisfied: cachetools<7,>=5.0.0 in /usr/local/lib/python3.11/dist-packages (from mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (5.5.2)\n",
            "Requirement already satisfied: click<9,>=7.0 in /usr/local/lib/python3.11/dist-packages (from mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (8.2.1)\n",
            "Requirement already satisfied: databricks-sdk<1,>=0.20.0 in /usr/local/lib/python3.11/dist-packages (from mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (0.57.0)\n",
            "Requirement already satisfied: fastapi<1 in /usr/local/lib/python3.11/dist-packages (from mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (0.115.13)\n",
            "Requirement already satisfied: gitpython<4,>=3.1.9 in /usr/local/lib/python3.11/dist-packages (from mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (3.1.44)\n",
            "Requirement already satisfied: importlib_metadata!=4.7.0,<9,>=3.7.0 in /usr/local/lib/python3.11/dist-packages (from mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (8.7.0)\n",
            "Requirement already satisfied: opentelemetry-api<3,>=1.9.0 in /usr/local/lib/python3.11/dist-packages (from mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (1.34.1)\n",
            "Requirement already satisfied: opentelemetry-sdk<3,>=1.9.0 in /usr/local/lib/python3.11/dist-packages (from mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (1.34.1)\n",
            "Requirement already satisfied: pydantic<3,>=1.10.8 in /usr/local/lib/python3.11/dist-packages (from mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (2.11.7)\n",
            "Requirement already satisfied: pyyaml<7,>=5.1 in /usr/local/lib/python3.11/dist-packages (from mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (6.0.2)\n",
            "Requirement already satisfied: sqlparse<1,>=0.4.0 in /usr/local/lib/python3.11/dist-packages (from mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (0.5.3)\n",
            "Requirement already satisfied: uvicorn<1 in /usr/local/lib/python3.11/dist-packages (from mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (0.34.3)\n",
            "Requirement already satisfied: Mako in /usr/lib/python3/dist-packages (from alembic!=1.10.0,<2->mlflow->google-meridian[and-cuda,colab,mlflow]) (1.1.3)\n",
            "Requirement already satisfied: wheel<1.0,>=0.23.0 in /usr/local/lib/python3.11/dist-packages (from astunparse>=1.6.0->tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (0.45.1)\n",
            "Requirement already satisfied: urllib3>=1.26.0 in /usr/local/lib/python3.11/dist-packages (from docker<8,>=4.0.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (2.4.0)\n",
            "Requirement already satisfied: blinker>=1.9.0 in /usr/local/lib/python3.11/dist-packages (from Flask<4->mlflow->google-meridian[and-cuda,colab,mlflow]) (1.9.0)\n",
            "Requirement already satisfied: itsdangerous>=2.2.0 in /usr/local/lib/python3.11/dist-packages (from Flask<4->mlflow->google-meridian[and-cuda,colab,mlflow]) (2.2.0)\n",
            "Requirement already satisfied: markupsafe>=2.1.1 in /usr/local/lib/python3.11/dist-packages (from Flask<4->mlflow->google-meridian[and-cuda,colab,mlflow]) (3.0.2)\n",
            "Requirement already satisfied: werkzeug>=3.1.0 in /usr/local/lib/python3.11/dist-packages (from Flask<4->mlflow->google-meridian[and-cuda,colab,mlflow]) (3.1.3)\n",
            "Requirement already satisfied: graphql-core<3.3,>=3.1 in /usr/local/lib/python3.11/dist-packages (from graphene<4->mlflow->google-meridian[and-cuda,colab,mlflow]) (3.2.6)\n",
            "Requirement already satisfied: graphql-relay<3.3,>=3.1 in /usr/local/lib/python3.11/dist-packages (from graphene<4->mlflow->google-meridian[and-cuda,colab,mlflow]) (3.2.0)\n",
            "Requirement already satisfied: attrs>=22.2.0 in /usr/local/lib/python3.11/dist-packages (from jsonschema>=3.0->altair>=5->google-meridian[and-cuda,colab,mlflow]) (25.3.0)\n",
            "Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /usr/local/lib/python3.11/dist-packages (from jsonschema>=3.0->altair>=5->google-meridian[and-cuda,colab,mlflow]) (2025.4.1)\n",
            "Requirement already satisfied: referencing>=0.28.4 in /usr/local/lib/python3.11/dist-packages (from jsonschema>=3.0->altair>=5->google-meridian[and-cuda,colab,mlflow]) (0.36.2)\n",
            "Requirement already satisfied: rpds-py>=0.7.1 in /usr/local/lib/python3.11/dist-packages (from jsonschema>=3.0->altair>=5->google-meridian[and-cuda,colab,mlflow]) (0.25.1)\n",
            "Requirement already satisfied: rich in /usr/local/lib/python3.11/dist-packages (from keras>=3.5.0->tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (13.9.4)\n",
            "Requirement already satisfied: namex in /usr/local/lib/python3.11/dist-packages (from keras>=3.5.0->tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (0.1.0)\n",
            "Requirement already satisfied: optree in /usr/local/lib/python3.11/dist-packages (from keras>=3.5.0->tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (0.16.0)\n",
            "Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=3.5->arviz->google-meridian[and-cuda,colab,mlflow]) (1.3.2)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=3.5->arviz->google-meridian[and-cuda,colab,mlflow]) (0.12.1)\n",
            "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=3.5->arviz->google-meridian[and-cuda,colab,mlflow]) (4.58.4)\n",
            "Requirement already satisfied: kiwisolver>=1.3.1 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=3.5->arviz->google-meridian[and-cuda,colab,mlflow]) (1.4.8)\n",
            "Requirement already satisfied: pillow>=8 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=3.5->arviz->google-meridian[and-cuda,colab,mlflow]) (11.2.1)\n",
            "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=3.5->arviz->google-meridian[and-cuda,colab,mlflow]) (3.2.3)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.11/dist-packages (from requests<3,>=2.21.0->tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (3.4.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.11/dist-packages (from requests<3,>=2.21.0->tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (3.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.11/dist-packages (from requests<3,>=2.21.0->tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (2025.6.15)\n",
            "Requirement already satisfied: threadpoolctl>=3.1.0 in /usr/local/lib/python3.11/dist-packages (from scikit-learn<2->mlflow->google-meridian[and-cuda,colab,mlflow]) (3.6.0)\n",
            "Requirement already satisfied: greenlet>=1 in /usr/local/lib/python3.11/dist-packages (from sqlalchemy<3,>=1.4.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (3.2.3)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.11/dist-packages (from tensorboard<2.19,>=2.18->tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (3.8.2)\n",
            "Requirement already satisfied: tensorboard-data-server<0.8.0,>=0.7.0 in /usr/local/lib/python3.11/dist-packages (from tensorboard<2.19,>=2.18->tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (0.7.2)\n",
            "Requirement already satisfied: google-auth~=2.0 in /usr/local/lib/python3.11/dist-packages (from databricks-sdk<1,>=0.20.0->mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (2.38.0)\n",
            "Requirement already satisfied: starlette<0.47.0,>=0.40.0 in /usr/local/lib/python3.11/dist-packages (from fastapi<1->mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (0.46.2)\n",
            "Requirement already satisfied: gitdb<5,>=4.0.1 in /usr/local/lib/python3.11/dist-packages (from gitpython<4,>=3.1.9->mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (4.0.12)\n",
            "Requirement already satisfied: zipp>=3.20 in /usr/local/lib/python3.11/dist-packages (from importlib_metadata!=4.7.0,<9,>=3.7.0->mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (3.23.0)\n",
            "Requirement already satisfied: opentelemetry-semantic-conventions==0.55b1 in /usr/local/lib/python3.11/dist-packages (from opentelemetry-sdk<3,>=1.9.0->mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (0.55b1)\n",
            "Requirement already satisfied: annotated-types>=0.6.0 in /usr/local/lib/python3.11/dist-packages (from pydantic<3,>=1.10.8->mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (0.7.0)\n",
            "Requirement already satisfied: pydantic-core==2.33.2 in /usr/local/lib/python3.11/dist-packages (from pydantic<3,>=1.10.8->mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (2.33.2)\n",
            "Requirement already satisfied: typing-inspection>=0.4.0 in /usr/local/lib/python3.11/dist-packages (from pydantic<3,>=1.10.8->mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (0.4.1)\n",
            "Requirement already satisfied: h11>=0.8 in /usr/local/lib/python3.11/dist-packages (from uvicorn<1->mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (0.16.0)\n",
            "Requirement already satisfied: markdown-it-py>=2.2.0 in /usr/local/lib/python3.11/dist-packages (from rich->keras>=3.5.0->tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (3.0.0)\n",
            "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /usr/local/lib/python3.11/dist-packages (from rich->keras>=3.5.0->tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (2.19.2)\n",
            "Requirement already satisfied: smmap<6,>=3.0.1 in /usr/local/lib/python3.11/dist-packages (from gitdb<5,>=4.0.1->gitpython<4,>=3.1.9->mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (5.0.2)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.11/dist-packages (from google-auth~=2.0->databricks-sdk<1,>=0.20.0->mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (0.4.2)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.11/dist-packages (from google-auth~=2.0->databricks-sdk<1,>=0.20.0->mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (4.9.1)\n",
            "Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.11/dist-packages (from markdown-it-py>=2.2.0->rich->keras>=3.5.0->tensorflow<2.19,>=2.18->google-meridian[and-cuda,colab,mlflow]) (0.1.2)\n",
            "Requirement already satisfied: anyio<5,>=3.6.2 in /usr/local/lib/python3.11/dist-packages (from starlette<0.47.0,>=0.40.0->fastapi<1->mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (4.9.0)\n",
            "Requirement already satisfied: sniffio>=1.1 in /usr/local/lib/python3.11/dist-packages (from anyio<5,>=3.6.2->starlette<0.47.0,>=0.40.0->fastapi<1->mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (1.3.1)\n",
            "Requirement already satisfied: pyasn1<0.7.0,>=0.6.1 in /usr/local/lib/python3.11/dist-packages (from pyasn1-modules>=0.2.1->google-auth~=2.0->databricks-sdk<1,>=0.20.0->mlflow-skinny==3.1.0->mlflow->google-meridian[and-cuda,colab,mlflow]) (0.6.1)\n"
          ]
        }
      ],
      "source": [
        "# Install meridian with mlflow: from PyPI @ latest release\n",
        "!pip install --upgrade google-meridian[colab,and-cuda,mlflow]\n",
        "\n",
        "# Install meridian with mlflow: from PyPI @ specific version\n",
        "# !pip install google-meridian[colab,and-cuda,mlflow]==1.1.3\n",
        "\n",
        "# Install meridian: from GitHub @HEAD\n",
        "# !pip install --upgrade \"google-meridian[colab,and-cuda,mlflow] @ git+https://github.com/google/meridian.git\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Fhwt1wzgLwpZ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b82accaf-dd64-42f2-c121-777bfe80b2fd"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Your runtime has 13.6 gigabytes of available RAM\n",
            "\n",
            "Num GPUs Available:  1\n",
            "Num CPUs Available:  1\n"
          ]
        }
      ],
      "source": [
        "# Install (latest) Numpy, Tensorflow and Meridian version\n",
        "import pandas as pd\n",
        "import tensorflow as tf\n",
        "import tensorflow_probability as tfp\n",
        "from meridian import constants\n",
        "from meridian.data import data_frame_input_data_builder\n",
        "from meridian.model import model\n",
        "from meridian.model import spec\n",
        "from meridian.model import prior_distribution\n",
        "\n",
        "# check if GPU is available\n",
        "from psutil import virtual_memory\n",
        "ram_gb = virtual_memory().total / 1e9\n",
        "print('Your runtime has {:.1f} gigabytes of available RAM\\n'.format(ram_gb))\n",
        "print(\"Num GPUs Available: \", len(tf.config.experimental.list_physical_devices('GPU')))\n",
        "print(\"Num CPUs Available: \", len(tf.config.experimental.list_physical_devices('CPU')))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HWLt_MdnvBrH"
      },
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kiM0UrN6qbIP"
      },
      "source": [
        "<a name=\"load-data\"></a>\n",
        "## Step 1: Load the data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z18Mo-22x0lY"
      },
      "source": [
        "Load the [simulated dataset in CSV format](https://github.com/google/meridian/blob/main/meridian/data/simulated_data/csv/geo_all_channels.csv) as follows."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tZd-ik8NbjK6"
      },
      "source": [
        "1\\. Read the data into a Pandas DataFrame."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7sV1ChiEYuyD"
      },
      "outputs": [],
      "source": [
        "df = pd.read_csv(\n",
        "    \"https://raw.githubusercontent.com/google/meridian/refs/heads/main/meridian/data/simulated_data/csv/geo_all_channels.csv\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8JBDZzl80BrY"
      },
      "source": [
        "2\\. Create a DataFrameInputDataBuilder instance.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4qdTSk4a0znn"
      },
      "outputs": [],
      "source": [
        "builder = data_frame_input_data_builder.DataFrameInputDataBuilder(\n",
        "    kpi_type='non_revenue'\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LNr75vQL1Zru"
      },
      "source": [
        "3\\. Offer the components to the builder. Note that the components may be offered all at once or piecewise.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "udaLGvwl1U8B"
      },
      "outputs": [],
      "source": [
        "builder = (\n",
        "    builder.with_kpi(df, kpi_col=\"conversions\")\n",
        "    .with_revenue_per_kpi(df, revenue_per_kpi_col=\"revenue_per_conversion\")\n",
        "    .with_population(df)\n",
        "    .with_controls(\n",
        "        df, control_cols=[\"sentiment_score_control\", \"competitor_sales_control\"]\n",
        "    )\n",
        ")\n",
        "\n",
        "channels = [\"Channel0\", \"Channel1\", \"Channel2\", \"Channel3\", \"Channel4\"]\n",
        "builder = builder.with_media(\n",
        "    df,\n",
        "    media_cols=[f\"{channel}_impression\" for channel in channels],\n",
        "    media_spend_cols=[f\"{channel}_spend\" for channel in channels],\n",
        "    media_channels=channels,\n",
        ")\n",
        "\n",
        "data = builder.build()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DlF5vs8vb8Wn"
      },
      "source": [
        "Note that the simulated data here does not contain reach and frequency. We recommend including reach and frequency data whenever they are available. For information about the advantages of utilizing reach and frequency, see [Bayesian Hierarchical Media Mix Model Incorporating Reach and Frequency Data](https://research.google/pubs/bayesian-hierarchical-media-mix-model-incorporating-reach-and-frequency-data/#:~:text=By%20incorporating%20R%26F%20into%20MMM,based%20on%20optimal%20frequency%20recommendations.). For code snippet for loading reach and frequency data, see [Load geo-level data with reach and frequency](https://developers.google.com/meridian/docs/user-guide/load-geo-data-with-rf)\n",
        "\n",
        "The documentation provides guidance for instances where reach and frequency data is accessible for specific channels. Additionally, for information about how to load other data types and formats, including data with reach and frequency, see [Supported data types and formats](https://developers.google.com/meridian/docs/user-guide/supported-data-types-formats)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FO6pDd6f2V1L"
      },
      "source": [
        "<a name=\"configure-model\"></a>\n",
        "## Step 2: Configure the model -- with MLflow Autolog!\n",
        "\n",
        "Note, this example focuses on MLflow autologging.\n",
        "\n",
        "For more information about configuring the parameters and using a customized model specification, such as setting different ROI priors for each media channel, see [Configure the model](https://developers.google.com/meridian/docs/user-guide/configure-model)."
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "1. Prepare the model specification."
      ],
      "metadata": {
        "id": "etXHONLdvBKz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "roi_mu = 0.2     # Mu for ROI prior for each media channel.\n",
        "roi_sigma = 0.9  # Sigma for ROI prior for each media channel.\n",
        "prior = prior_distribution.PriorDistribution(\n",
        "    roi_m=tfp.distributions.LogNormal(roi_mu, roi_sigma, name=constants.ROI_M)\n",
        ")\n",
        "model_spec = spec.ModelSpec(prior=prior)"
      ],
      "metadata": {
        "id": "qc7VV2adu9Kr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "2. Enable MLflow autologging by calling `autolog()`. This enables automatic logging of parameters, metrics, and models by patching relevant functions. This only needs to be called once prior to any model runs."
      ],
      "metadata": {
        "id": "ybqyMpQWwrol"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from meridian.mlflow import autolog\n",
        "\n",
        "autolog.autolog(log_metrics=True)  # Metric logging is not enabled by default."
      ],
      "metadata": {
        "id": "6wWOqK2vxD5e"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a_mQI7HzxxK4"
      },
      "source": [
        "3. In the context of a new MLflow run started with `mlflow.start_run()`, initialize the `Meridian` class and invoke `sample_prior()` and `sample_posterior()` to obtain samples from the prior and posterior distributions of model parameters. If you are using the T4 GPU runtime this step may take about 10 minutes for the provided data set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8XNDd7HX1qTn"
      },
      "outputs": [],
      "source": [
        "import mlflow\n",
        "\n",
        "with mlflow.start_run(run_name=\"my-run\"):\n",
        "  mmm = model.Meridian(input_data=data, model_spec=model_spec)\n",
        "  mmm.sample_prior(500)\n",
        "  mmm.sample_posterior(n_chains=7, n_adapt=500, n_burnin=500, n_keep=1000, seed=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "4. After the run completes, you can retrieve run results using the MLflow client."
      ],
      "metadata": {
        "id": "04Ec808ALZDf"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nZtQ1IjPRNbV",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "da7787d8-3f90-432d-8b8c-d01257e6e44b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "<Run: data=<RunData: metrics={'MAPE': 0.2559046149253845,\n",
            " 'R_Squared': 0.7730023860931396,\n",
            " 'wMAPE': 0.20020413398742676}, params={'arviz_version': '0.21.0',\n",
            " 'meridian_version': '1.1.3',\n",
            " 'prior.alpha_m': 'tfp.distributions.Uniform(\"alpha_m\", batch_shape=[], '\n",
            "                  'event_shape=[], dtype=float32)',\n",
            " 'prior.alpha_om': 'tfp.distributions.Uniform(\"alpha_om\", batch_shape=[], '\n",
            "                   'event_shape=[], dtype=float32)',\n",
            " 'prior.alpha_orf': 'tfp.distributions.Uniform(\"alpha_orf\", batch_shape=[], '\n",
            "                    'event_shape=[], dtype=float32)',\n",
            " 'prior.alpha_rf': 'tfp.distributions.Uniform(\"alpha_rf\", batch_shape=[], '\n",
            "                   'event_shape=[], dtype=float32)',\n",
            " 'prior.beta_m': 'tfp.distributions.HalfNormal(\"beta_m\", batch_shape=[], '\n",
            "                 'event_shape=[], dtype=float32)',\n",
            " 'prior.beta_om': 'tfp.distributions.HalfNormal(\"beta_om\", batch_shape=[], '\n",
            "                  'event_shape=[], dtype=float32)',\n",
            " 'prior.beta_orf': 'tfp.distributions.HalfNormal(\"beta_orf\", batch_shape=[], '\n",
            "                   'event_shape=[], dtype=float32)',\n",
            " 'prior.beta_rf': 'tfp.distributions.HalfNormal(\"beta_rf\", batch_shape=[], '\n",
            "                  'event_shape=[], dtype=float32)',\n",
            " 'prior.contribution_m': 'tfp.distributions.Beta(\"contribution_m\", '\n",
            "                         'batch_shape=[], event_shape=[], dtype=float32)',\n",
            " 'prior.contribution_n': 'tfp.distributions.TruncatedNormal(\"contribution_n\", '\n",
            "                         'batch_shape=[], event_shape=[], dtype=float32)',\n",
            " 'prior.contribution_om': 'tfp.distributions.Beta(\"contribution_om\", '\n",
            "                          'batch_shape=[], event_shape=[], dtype=float32)',\n",
            " 'prior.contribution_orf': 'tfp.distributions.Beta(\"contribution_orf\", '\n",
            "                           'batch_shape=[], event_shape=[], dtype=float32)',\n",
            " 'prior.contribution_rf': 'tfp.distributions.Beta(\"contribution_rf\", '\n",
            "                          'batch_shape=[], event_shape=[], dtype=float32)',\n",
            " 'prior.ec_m': 'tfp.distributions.TruncatedNormal(\"ec_m\", batch_shape=[], '\n",
            "               'event_shape=[], dtype=float32)',\n",
            " 'prior.ec_om': 'tfp.distributions.TruncatedNormal(\"ec_om\", batch_shape=[], '\n",
            "                'event_shape=[], dtype=float32)',\n",
            " 'prior.ec_orf': 'tfp.distributions.TransformedDistribution(\"ec_orf\", '\n",
            "                 'batch_shape=[], event_shape=[], dtype=float32)',\n",
            " 'prior.ec_rf': 'tfp.distributions.TransformedDistribution(\"ec_rf\", '\n",
            "                'batch_shape=[], event_shape=[], dtype=float32)',\n",
            " 'prior.eta_m': 'tfp.distributions.HalfNormal(\"eta_m\", batch_shape=[], '\n",
            "                'event_shape=[], dtype=float32)',\n",
            " 'prior.eta_om': 'tfp.distributions.HalfNormal(\"eta_om\", batch_shape=[], '\n",
            "                 'event_shape=[], dtype=float32)',\n",
            " 'prior.eta_orf': 'tfp.distributions.HalfNormal(\"eta_orf\", batch_shape=[], '\n",
            "                  'event_shape=[], dtype=float32)',\n",
            " 'prior.eta_rf': 'tfp.distributions.HalfNormal(\"eta_rf\", batch_shape=[], '\n",
            "                 'event_shape=[], dtype=float32)',\n",
            " 'prior.gamma_c': 'tfp.distributions.Normal(\"gamma_c\", batch_shape=[], '\n",
            "                  'event_shape=[], dtype=float32)',\n",
            " 'prior.gamma_n': 'tfp.distributions.Normal(\"gamma_n\", batch_shape=[], '\n",
            "                  'event_shape=[], dtype=float32)',\n",
            " 'prior.knot_values': 'tfp.distributions.Normal(\"knot_values\", batch_shape=[], '\n",
            "                      'event_shape=[], dtype=float32)',\n",
            " 'prior.mroi_m': 'tfp.distributions.LogNormal(\"mroi_m\", batch_shape=[], '\n",
            "                 'event_shape=[], dtype=float32)',\n",
            " 'prior.mroi_rf': 'tfp.distributions.LogNormal(\"mroi_rf\", batch_shape=[], '\n",
            "                  'event_shape=[], dtype=float32)',\n",
            " 'prior.roi_m': 'tfp.distributions.LogNormal(\"roi_m\", batch_shape=[], '\n",
            "                'event_shape=[], dtype=float32)',\n",
            " 'prior.roi_rf': 'tfp.distributions.LogNormal(\"roi_rf\", batch_shape=[], '\n",
            "                 'event_shape=[], dtype=float32)',\n",
            " 'prior.sigma': 'tfp.distributions.HalfNormal(\"sigma\", batch_shape=[], '\n",
            "                'event_shape=[], dtype=float32)',\n",
            " 'prior.slope_m': 'tfp.distributions.Deterministic(\"slope_m\", batch_shape=[], '\n",
            "                  'event_shape=[], dtype=float32)',\n",
            " 'prior.slope_om': 'tfp.distributions.Deterministic(\"slope_om\", '\n",
            "                   'batch_shape=[], event_shape=[], dtype=float32)',\n",
            " 'prior.slope_orf': 'tfp.distributions.LogNormal(\"slope_orf\", batch_shape=[], '\n",
            "                    'event_shape=[], dtype=float32)',\n",
            " 'prior.slope_rf': 'tfp.distributions.LogNormal(\"slope_rf\", batch_shape=[], '\n",
            "                   'event_shape=[], dtype=float32)',\n",
            " 'prior.tau_g_excl_baseline': 'tfp.distributions.Normal(\"tau_g_excl_baseline\", '\n",
            "                              'batch_shape=[], event_shape=[], dtype=float32)',\n",
            " 'prior.xi_c': 'tfp.distributions.HalfNormal(\"xi_c\", batch_shape=[], '\n",
            "               'event_shape=[], dtype=float32)',\n",
            " 'prior.xi_n': 'tfp.distributions.HalfNormal(\"xi_n\", batch_shape=[], '\n",
            "               'event_shape=[], dtype=float32)',\n",
            " 'sample_posterior.dual_averaging_kwargs': 'None',\n",
            " 'sample_posterior.init_step_size': 'None',\n",
            " 'sample_posterior.max_energy_diff': '500.0',\n",
            " 'sample_posterior.max_tree_depth': '10',\n",
            " 'sample_posterior.n_adapt': '500',\n",
            " 'sample_posterior.n_burnin': '500',\n",
            " 'sample_posterior.n_chains': '7',\n",
            " 'sample_posterior.n_keep': '1000',\n",
            " 'sample_posterior.parallel_iterations': '10',\n",
            " 'sample_posterior.seed': '1',\n",
            " 'sample_posterior.unrolled_leapfrog_steps': '1',\n",
            " 'sample_prior.n_draws': '500',\n",
            " 'sample_prior.seed': 'None',\n",
            " 'spec.baseline_geo': 'None',\n",
            " 'spec.control_population_scaling_id': 'None',\n",
            " 'spec.hill_before_adstock': 'False',\n",
            " 'spec.holdout_id': 'None',\n",
            " 'spec.knots': 'None',\n",
            " 'spec.max_lag': '8',\n",
            " 'spec.media_effects_dist': 'log_normal',\n",
            " 'spec.media_prior_type': 'None',\n",
            " 'spec.non_media_baseline_values': 'None',\n",
            " 'spec.non_media_population_scaling_id': 'None',\n",
            " 'spec.non_media_treatments_prior_type': 'contribution',\n",
            " 'spec.organic_media_prior_type': 'contribution',\n",
            " 'spec.organic_rf_prior_type': 'contribution',\n",
            " 'spec.paid_media_prior_type': 'None',\n",
            " 'spec.rf_prior_type': 'None',\n",
            " 'spec.rf_roi_calibration_period': 'None',\n",
            " 'spec.roi_calibration_period': 'None',\n",
            " 'spec.unique_sigma_for_each_geo': 'False'}, tags={'mlflow.runName': 'my-run',\n",
            " 'mlflow.source.name': '/usr/local/lib/python3.11/dist-packages/colab_kernel_launcher.py',\n",
            " 'mlflow.source.type': 'LOCAL',\n",
            " 'mlflow.user': 'root'}>, info=<RunInfo: artifact_uri='file:///content/mlruns/0/87edf3efd4ab4ddc96c176d4052ac08b/artifacts', end_time=1750807630540, experiment_id='0', lifecycle_stage='active', run_id='87edf3efd4ab4ddc96c176d4052ac08b', run_name='my-run', start_time=1750807234088, status='FINISHED', user_id='root'>, inputs=<RunInputs: dataset_inputs=[], model_inputs=[]>, outputs=<RunOutputs: model_outputs=[]>>\n"
          ]
        }
      ],
      "source": [
        "client = mlflow.tracking.MlflowClient()\n",
        "experiment_id = \"0\"\n",
        "runs = client.search_runs(\n",
        "    experiment_id,\n",
        "    max_results=1000,\n",
        "    filter_string=f\"attributes.run_name = 'my-run'\"\n",
        ")\n",
        "if runs:\n",
        "  print(runs[0])\n",
        "else:\n",
        "  print(\"No runs found.\")\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "language_info": {
      "name": "python"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
